import numpy as np
import matplotlib.pyplot as plt
from itertools import product, combinations

times = np.loadtxt('time3.5.txt')

idx = 40
t = np.linspace(0,1000,25001)
a = np.loadtxt('Li%s.txt'%(str(idx)))

fig = plt.figure(figsize=(8,6))
ax = plt.axes(projection='3d')
ax.set_box_aspect([1,1,1])

ax.plot3D(a[:,0],a[:,1],a[:,2],lw=0.1,zorder=0)
ax.scatter(a[0,0],a[0,1],a[0,2],color='blue',label='_nolegend_')

ax.plot3D(a[:,0],a[:,1],a[:,2],lw=1,color='green',zorder=1)

for i in range(len(times)):
#    r = globals()['r%s' % i]
    r = ((t[:] >= times[i]-0.1) & (t[:] <= times[i]+0.1))
    ax.plot3D(a[r,0],a[r,1],a[r,2],lw=2,color='red',zorder=1)

a = 2.0104010492000000e+01/2
# draw cube
r = [a, a*2]
for s, e in combinations(np.array(list(product(r, r, r))), 2):
    if np.sum(np.abs(s-e)) == r[1]-r[0]:
        ax.plot3D(*zip(s, e), color="black",lw=0.5)

'''
ax.set_xticks([5,10,15,20])
ax.set_yticks([5,10,15,20])
ax.set_zticks([0,5,10,15,20])
ax.axes.set_xlim3d(left=5, right=20)
ax.axes.set_ylim3d(bottom=5, top=20)
ax.axes.set_zlim3d(bottom=0, top=20)
'''
ax.set_axis_off()
ax.xaxis.set_rotate_label(False)
ax.yaxis.set_rotate_label(False)
ax.zaxis.set_rotate_label(False)
ax.set_xlabel(r'x ($\AA$)',fontsize=16,labelpad=6,rotation=0)
ax.set_ylabel(r'y ($\AA$)',fontsize=16,labelpad=6,rotation=0)
ax.set_zlabel(r'z ($\AA$)',fontsize=16,labelpad=6,rotation=90)
ax.grid(False)
ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.tick_params(direction='in',labelsize=16,pad=2)
#ax.view_init(8, -60)
ax.view_init(9, -45)
#plt.savefig('TrajLi%s_range.pdf'%(str(idx)),bbox_inches='tight',format='pdf')
plt.savefig('TrajLi%s_range_Dec23.pdf'%(str(idx)),format='pdf')
plt.savefig('TrajLi%s_range_Dec23.eps'%(str(idx)),format='eps')
plt.savefig('TrajLi%s_range_Dec23.png'%(str(idx)),format='png',dpi=900)
plt.show()
